# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import re
import dashscope
import unicodedata
from typing import Union, Optional, List, Dict
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.embeddings import SparkLLMTextEmbeddings
from langchain_community.embeddings import DashScopeEmbeddings
from common.enums.models import ModelsEnum
from common.chain.chain_server import ChatUtils
from common.cache.pool_cache import KeyPoolCache
from common.models.dev import DevModelsModel


class VectorService:
    # ====== 以下使用预编译正则提升性能 ======
    # HTML/MD图片的链接
    _ALT_PATTERN = re.compile(r'alt=["\'](.*?)["\']')
    _HTML_IMG_PATTERN = re.compile(r"<img[^>]+>")
    _MD_IMG_PATTERN = re.compile(r"!\[.*?]\(.*?\)")
    _MD_DESC_PATTERN = re.compile(r'!\[(.*?)]')
    # 空白字符和标点符号
    _WHITESPACE_PATTERN = re.compile(r"\s+")
    _PARAGRAPHS_PATTERN = re.compile(r"\n{2,}")
    _PUNCTUATION_PATTERN = re.compile(r"([!?,.;:])\1+")
    _LINE_WHITESPACE_PATTERN = re.compile(r"^\s+|\s+$", flags=re.MULTILINE)
    # 规范化引号和破折号
    _QUOTES_PATTERN1 = re.compile(r"[\u201c\u201d\u0022]")
    _QUOTES_PATTERN2 = re.compile(r"[\u2018\u2019\u0027]")
    _DASH_PATTERN = re.compile(r"[—–]")
    # 中文特定清洗规则
    _CHINESE_PATTERN = re.compile(r"[\u4e00-\u9fff]")
    _FULLWIDTH_PATTERN = re.compile(r"[\uFF01-\uFF5E\u3000]")
    _CN_PUNCTUATION_PATTERN = re.compile(r"([。，！？；：、])\1+")
    _CN_QUOTES_PATTERN = re.compile(r"[\u300c\u300d\u300e\u300f]")
    # ====== 以上使用预编译正则提升性能 ======

    def __init__(self):
        self.usage_list = []
        self.usage = {
            "total_tokens": 0,
            "prompt_tokens": 0,
            "completion_tokens": 0
        }

    async def to_embed(self, models: Union[str, dict], text: Union[str, List[str]]):
        try:
            if isinstance(models, str):
                models = await DevModelsModel.check_models(models, ModelsEnum.TYPE_VECTOR)
        except Exception as e:
            raise Exception(*self._format_error(str(e)))

        try:
            poolKeyCache = KeyPoolCache(scene=ModelsEnum.TYPE_VECTOR, channel=models["mark"])
            keys: Dict[str, Union[int, str]] = await poolKeyCache.get_key()
            if keys is None:
                raise Exception("The api_key client")
        except Exception as e:
            raise Exception(*self._format_error(str(e)))

        apiKey: str = keys.get("api_key")
        modelName: str = models.get("model")
        agencyApi: str = models.get("config").get("agency_api")

        try:
            if models.get("channel") == "openai":
                return await self.openai(text, modelName, apiKey, agencyApi)
            elif models.get("channel") == "zhipu":
                return await self.zhipu(text, modelName, apiKey, agencyApi)
            elif models.get("channel") == "xunfei":
                return await self.xunfei(text, keys, agencyApi)
            elif models.get("channel") == "qwen":
                return await self.qwen(text, modelName, apiKey, agencyApi)
        except Exception as e:
            if e.args[0] in [402, 403, 404]:  # 都属于key失效
                await poolKeyCache.down_key(e.args[1])
            raise e

    async def openai(self, text: Union[str, List[str]], model: str, api_key: str, base_url: Optional[str] = None):
        """
        OpenAi: https://api.openai.com

        Args:
            text (Union[str, List[str]]): 文本内容。
            model (str): 模型。
            api_key (str): apiKey。
            base_url (Optional[str]): baseUrl。

        Author:
            zero
        """
        try:
            api_base = base_url or "https://api.openai.com/v1"
            os.environ["OPENAI_API_BASE"] = api_base
            os.environ["OPENAI_API_KEY"] = api_key

            embeddings = OpenAIEmbeddings(model=model)

            self._usage_tokens(text)
            if isinstance(text, str):
                text = self._cleaning_text(text)
                return await embeddings.aembed_query(text)
            else:
                texts = [self._cleaning_text(t) for t in text]
                return await embeddings.aembed_documents(texts)
        except Exception as e:
            raise Exception(*self._format_error(str(e)))

    async def zhipu(self, text: Union[str, List[str]], model: str, api_key: str, api_base: Optional[str] = None):
        """
        智谱清言: https://open.bigmodel.cn/dev/api/vector/embedding-3
        """
        try:
            if api_base:
                return await self.openai(text, model, api_key, api_base)

            embeddings = ZhipuAIEmbeddings(
                api_key=api_key,
                model=model
            )

            self._usage_tokens(text)
            if isinstance(text, str):
                text = self._cleaning_text(text)
                return await embeddings.aembed_query(text)
            else:
                texts = [self._cleaning_text(t) for t in text]
                return await embeddings.aembed_documents(texts)
        except Exception as e:
            raise Exception(*self._format_error(str(e)))

    async def xunfei(self, text: Union[str, List[str]], keys: dict, base_url: Optional[str] = None):
        """
        讯飞星火: https://www.xfyun.cn/doc/spark/Embedding_api.html

        Args:
            text (Union[str, List[str]]): 文本内容。
            keys (dict): apiKey。
            base_url (Optional[str]): baseUrl。

        Author:
            zero
        """
        try:
            api_base = base_url or "https://emb-cn-huabei-1.xf-yun.com/"
            api_base = api_base if api_base.endswith("/") else api_base + "/"
            embeddings = SparkLLMTextEmbeddings(
                base_url=api_base,
                app_id=keys.get("app_id"),
                api_key=keys.get("api_key"),
                api_secret=keys.get("api_secret")
            )

            self._usage_tokens(text)
            if isinstance(text, str):
                text = self._cleaning_text(text)
                result = await embeddings.aembed_query(text)
            else:
                texts = [self._cleaning_text(t) for t in text]
                result = await embeddings.aembed_documents(texts)

            if not isinstance(result, list):
                raise Exception("向量化解析失败: " + str(result))
            return result
        except Exception as e:
            raise Exception(*self._format_error(str(e)))

    async def qwen(self, text: Union[str, List[str]], model: str, api_key: str, base_url: Optional[str] = None):
        """
        通义千问: https://help.aliyun.com/zh/model-studio/developer-reference/text-embedding-synchronous-api

        Args:
            text (Union[str, List[str]]): 文本内容。
            model (str): 模型名称。
            api_key (str): apiKey。
            base_url (Optional[str]): baseUrl。

        Author:
            zero
        """
        try:
            if base_url:
                dashscope.base_http_api_url = base_url

            embeddings = DashScopeEmbeddings(
                model=model,
                dashscope_api_key=api_key
            )

            self._usage_tokens(text)
            if isinstance(text, str):
                text = self._cleaning_text(text)
                return await embeddings.aembed_query(text)
            else:
                texts = [self._cleaning_text(t) for t in text]
                return await embeddings.aembed_documents(texts)
        except Exception as e:
            raise Exception(*self._format_error(str(e)))

    def _usage_tokens(self, text: Union[str, List[str]]) -> dict[str, int] | None:
        """
        获取消费Tokens信息

        Args:
            text (Union[str, List[str]]): 输入文本

        Author:
            zero

        Returns:
            tokens (dict)
        """
        if isinstance(text, str):
            tokens: int = ChatUtils.compute_tokens(text)
            self.usage["total_tokens"] = tokens
            self.usage["prompt_tokens"] = tokens
            return self.usage
        elif isinstance(text, list):
            usage_arr: List[Dict[str, int]] = []
            total_text: str = ""
            for t in text:
                total_text += t
                tokens: int = ChatUtils.compute_tokens(t)
                usage_arr.append({
                    "total_tokens": tokens,
                    "prompt_tokens": tokens,
                    "completion_tokens": 0
                })

            tokens: int = ChatUtils.compute_tokens(total_text)
            self.usage["total_tokens"] = tokens
            self.usage["prompt_tokens"] = tokens
            self.usage_list = usage_arr
            return self.usage

        return None

    @classmethod
    def _format_error(cls, error) -> tuple:
        """
        错误信息

        Args:
            error (str): 错误信息

        Author:
            zero

        Returns:
            error (tuple)
        """
        error = error.strip()
        if error.startswith("Error code: 400"):
            return 400, "Vector: 模型不可用，请管理员检查模型代码。"
        elif error.startswith("Connection error"):
            return 401, "Vector: 请求已超时，请管理员检查代理地址。"
        elif error.startswith("Error code: 401"):
            return 402, "Vector: 无效的密钥，请管理员检查密钥配置。"
        elif error.startswith("The api_key client"):
            return 403, "Vector: api_key无效，请管理员检查密钥配置。"
        elif error.startswith("Error code: 429"):
            return 404, "Vector: api_key失效，KEY账户已欠费。"
        else:
            return 500, error

    @classmethod
    def _cleaning_text(cls, text) -> str:
        """
        对文本进行通用预处理, 为向量化做准备

        Args:
            text (str): 需要清洗的原始文本

        Returns:
            str: 清洗后的文本

        Author:
            zero
        """
        if not text or not isinstance(text, str):
            return ""

        # 检查是否包含中文字符 (提前检测,避免后续不必要的处理)
        has_chinese = bool(cls._CHINESE_PATTERN.search(text))

        # 1. 标准化Unicode字符
        text = unicodedata.normalize("NFKC", text)

        # 2. 替换HTML图片标签和Markdown图片
        text = cls._HTML_IMG_PATTERN.sub(cls._extract_alt_text, text)
        text = cls._MD_IMG_PATTERN.sub(cls._extract_md_image_desc, text)

        # 3. 批量处理空白字符和标点符号
        text = cls._WHITESPACE_PATTERN.sub(" ", text)
        text = cls._PARAGRAPHS_PATTERN.sub('\n\n', text)
        text = cls._PUNCTUATION_PATTERN.sub(r"\1", text)
        text = cls._LINE_WHITESPACE_PATTERN.sub("", text)

        # 4. 规范化引号和破折号
        text = cls._QUOTES_PATTERN1.sub('"', text)
        text = cls._QUOTES_PATTERN2.sub("'", text)
        text = cls._DASH_PATTERN.sub("-", text)

        # 5. 中文特定清洗规则
        if has_chinese:
            text = cls._FULLWIDTH_PATTERN.sub(cls._full_to_half, text)
            text = cls._CN_PUNCTUATION_PATTERN.sub(r"\1", text)
            text = cls._CN_QUOTES_PATTERN.sub('"', text)

        # 6. 去除文本首尾的空白（使用字符串方法而非正则）
        return text.strip()

    @staticmethod
    def _extract_alt_text(match):
        """ 提取HTML图片标签中的alt文本 """
        alt_match = VectorService._ALT_PATTERN.search(match.group(0))
        return " " + (alt_match.group(1) if alt_match else "") + " "

    @staticmethod
    def _extract_md_image_desc(match):
        """ 提取Markdown图片链接中的描述文本 """
        desc_match = VectorService._MD_DESC_PATTERN.search(match.group(0))
        return " " + (desc_match.group(1) if desc_match and desc_match.group(1) else "") + " "

    @staticmethod
    def _full_to_half(match):
        """ 全角字符转半角字符 """
        char = match.group(0)
        code = ord(char)
        if code == 0x3000:  # 全角空格
            return " "
        elif 0xFF01 <= code <= 0xFF5E:
            return chr(code - 0xFEE0)
        return char
